package com.example.mcpclient.repository;

import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.filter.Filter;
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Repository;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author xuguang
 */
@Repository
public class ElasticsearchMemoryRepository implements ChatMemoryRepository {

    @Autowired
    private VectorStore vectorStore;

    Map<String, List<Message>> chatMemoryStore = new ConcurrentHashMap<>();

    @Override
    public List<String> findConversationIds() {
        return new ArrayList<>(this.chatMemoryStore.keySet());
    }

    @Override
    public List<Message> findByConversationId(String conversationId) {
        SearchRequest searchRequest = SearchRequest.builder()
                .filterExpression(toFilter(conversationId))
                .build();
        List<Document> documents = vectorStore.similaritySearch(searchRequest);
        return documents.stream().map(DocumentMessageConvert::toMessage).toList();
    }

    @Override
    public void saveAll(String conversationId, List<Message> messages) {
        List<Document> documents = DocumentMessageConvert.toDocuments(conversationId, messages);
        vectorStore.add(documents);
    }

    @Override
    public void deleteByConversationId(String conversationId) {
        vectorStore.delete(toFilter(conversationId));
    }

    /**
     * 转换为filter
     */
    private Filter.Expression toFilter(String conversationId) {
        FilterExpressionBuilder builder = new FilterExpressionBuilder();
        return builder.eq("conversationId", conversationId).build();
    }


}
